/*
 * Stage.cpp
 *
 *  Created on: Jun 9, 2010
 *      Author: bernardo
 */
#include <algorithm>
#include <iostream>
#include <cmath>
#include <vector>

#ifndef PROBDENSITY_H_
#define PROBDENSITY_H_

template<size_t nComp>
class ProbDens { //Probability density field with nComp components
private:

public:
	class CompIterator;
	typedef CompIterator iterator;
	typedef const CompIterator const_iterator;

	explicit ProbDens(size_t nX, double x0, double x1);

	explicit ProbDens(double dX);

	explicit ProbDens() {}

	~ProbDens() {}

	void clear(size_t nX, double x0, double x1);

	void clear(double dX);

	inline double getX(size_t i) const {
		return x0 + (i + 0.5) * dX;
	}

	inline double getDX() const {
		return dX;
	}

	inline double operator[](size_t i) const;

	inline size_t getNX() const {
		return dens.size()/nComp;
	}

	inline iterator begin(size_t comp) {
		return iterator(dens.begin() + comp);
	}

	inline iterator end(size_t comp) {
		return begin(comp) + getNX();
	}

	inline iterator iter(size_t comp, size_t i) {
		return begin(comp) +  i;
	}

	inline const_iterator iter(size_t comp, size_t i) const {
		return const_cast<ProbDens<nComp>* >(this)->iter(comp,i);
	}

	inline iterator iter(size_t comp,double x) {
		return iter(comp, size_t((x-x0)/dX));
	}

	inline const_iterator iter(size_t comp, double x) const {
		return iter(comp, size_t((x-x0)/dX));
	}

	void add(double x, double dens);

	void deltaFunction(double x,double *prob) ;

	void saveProbDens(std::ostream& out, double time) const;

private:
	double x0,x1,dX;
	std::vector<double> dens;
	//prevent undesired copy and creation (not to be implemented)
	ProbDens(const ProbDens&);
	ProbDens& operator=(const ProbDens&);
};

template<size_t nComp>
class  ProbDens<nComp>::CompIterator {
	//Iterates within one component
	public:
		explicit CompIterator(std::vector<double>::iterator p) : p(p) {};
		//If necessary, operator += etc must be redefined
		inline CompIterator operator+(int i) const {return CompIterator(p+i*nComp);}
		inline CompIterator& operator+=(int i) {return *this=*this+i;}
		inline CompIterator& operator++() { return *this+=1; }
		inline CompIterator operator++(int) {CompIterator cI(*this); ++(*this); return cI; }
		inline int operator-(CompIterator cI) const {return (p-cI.p)/nComp;}
		inline bool operator!=(CompIterator cI) const {return p!=cI.p;}
		inline bool operator<(CompIterator cI) const {return p<cI.p;}
		inline bool operator>=(CompIterator cI) const {return p>=cI.p;}
		inline double& operator*() {return *p;}
		inline const double operator*() const {return *p;}
	private:
		typename std::vector<double>::iterator p;
	};

template<size_t nComp>
ProbDens<nComp>::ProbDens(size_t nX, double x0, double x1) :
	x0(x0), x1(x1), dX((x1-x0)/nX), dens(nX*nComp,0) {
}

template<size_t nComp>
ProbDens<nComp>::ProbDens(double dX) :
	dX(dX) {
}

template<size_t nComp>
inline double ProbDens<nComp>::operator[](size_t i) const {
	double prob=*iter(0,i);
	for (size_t c=0; c!=nComp; c++) {
		prob+=*iter(0,i);
	}
	return prob;
}


template<size_t nComp>
void ProbDens<nComp>::clear(size_t nX, double x0, double x1) {
	this->x0=x0;
	this->x1=x1;
	dX=(x1-x0)/nX;
	dens.resize(nX*nComp);
	fill(dens.begin(),dens.end(),0);
}

template<size_t nComp>
void ProbDens<nComp>::clear(double dX) {
	this->dX=dX;
	dens.clear();
}

template<size_t nComp>
void ProbDens<nComp>::add(double x, double prob) {
	if (getNX()==0) {
		x1=x0=x-dX/2;
	}
	if (iter(0,x)<begin(0)) {
		std::cerr << "Error in probDensity: " << x << "\n";
		exit(1);
	}
	if (iter(0,x)>=end(0)) {
		//std::cout << "nx:" << getNX() << "\t"<< x0 << "\t" << x1 << "\n";
		dens.resize((iter(0,x)+1-begin(0))*nComp,0);
		x1 = x0 + getNX()*dX;
	}
	for (size_t i=0; i!=nComp; i++) {
		*iter(i,x) += prob/dX/nComp;
	}
}

template<size_t nComp>
void ProbDens<nComp>::deltaFunction(double x,double* prob) {
	for (size_t j=0; j!=nComp; j++)
		*iter(j,x) = *(prob++)/dX;
}

template<size_t nComp>
void ProbDens<nComp>::saveProbDens(std::ostream& out, double time) const {
	out << "#time\tphi";
	for (size_t c = 0; c!= nComp; c++)
		out << "\tdens" << c;
	out << "\n";
	if (nComp==1) {
		std::cout << "\n" << x0 << "\t" << x1 << "\t" << getNX() << "\n" ;
	}
	for (size_t i = 0; i != getNX(); i++) {
		out << time << "\t" << getX(i);
		for (size_t c = 0; c != nComp; c++) out  << "\t" << *iter(c,i);
		out << "\n";
	}
}


#endif /* PROBDENSITY_H_ */
